from langchain.prompts import PromptTemplate
from pydantic import BaseModel, Field
from datetime import datetime
import json
from LLM_Prompt.siliconflow import SiliconFlow


# Step 1: 使用 SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()  # 替换成你实际的 SiliconFlow 模型接口

classification_prompt = PromptTemplate(
    input_variables=["question"],
    template="""
**任务说明**  
你的任务是判断用户输入的问题是否与法律、法规、规章、政策或合规事项、航海操作技术范畴、船舶垃圾的分类、上海地理范围相关。请只根据问题内容给出判断，并返回规定 JSON 格式的结果。

# 输出要求

请仅返回如下 JSON 结构：

```json
{{
  "related_to_law": true 或 false,
  "reason": "简要说明判断理由"
}}
```

## 要求

1. 只返回以上 JSON 格式结果，不要输出其他内容。
2. 如涉及法律、法规、行政规章、合规、政策条文、权限、义务、处罚、司法解释、航海操作技术范畴、船舶垃圾的分类、上海地理范围等任意一项，均属相关。
3. 有时间范围、具体船名、卡口名称中任意一项，则一律视为不相关。
4. 涉及黄浦江的问题，均属相关。

输入问题：{question}

输出  .
严格按要求输出。
"""
)

# Step 4: Pydantic 数据验证器
class ExtractedData(BaseModel):
    boat_name: str = Field(..., description="船的名字")

    @staticmethod
    def validate_output(response: str):
        try:
            data = ExtractedData.model_validate_json(response)  # 替换 parse_raw
            return data.model_dump()  # 替换 dict
        except Exception as e:
            raise ValueError(f"数据解析错误: {e}")


# 清理模型输出工具
def clean_model_output(raw_output: str) -> str:
    """
    清理大模型的输出，去除额外的 Markdown 等格式符号，确保输出是合法 JSON。
    """
    # 去掉 Markdown 格式的代码块标记（如 ```json ... ```)
    if raw_output.startswith("```json"):
        raw_output = raw_output.strip("```json").strip("`\n")
    elif raw_output.startswith("```"):
        raw_output = raw_output.strip("```").strip("`\n")

    try:
        # 确保最终输出为合法的 JSON 格式（如有多余的空格，再清理一下）
        raw_output = json.loads(raw_output)  # 测试是否是有效 JSON
    except json.JSONDecodeError as e:
        raise ValueError(f"清理后的输出仍不是有效 JSON: {e}")

    return raw_output


# Step 5: 链式操作
classification_chain = classification_prompt | llm


def classify_and_extract(question):

    classification_extract_json = classification_chain.invoke({"question": question})

    # 清理模型输出
    cleaned_result = clean_model_output(classification_extract_json)

    return  cleaned_result


# 测试代码

test_queries = [
        "莲花号前天的航道轨迹",
        "莲花号2025年1月20号的航道轨迹",
        "船舶 ABC123 2025年2月20日的行驶轨迹？",
        "船舶 XYZ789 昨天经过那些卡口？",
        "船舶 DEF456 本周是否有违规行为（预警）？",
        "船舶 GHI789 上周去过哪些码头？去过哪些区域？途径那些船闸？途径那些航道？",
        "卡口 南京长江大桥 2025年2月20日到22日有哪些船舶经过？有哪些违规船舶？有哪些集装箱船舶经过等",
        "航道 长江主航道 当前有哪些船舶？有哪些违规船舶？有哪些类别船舶？", # 这个问题无法很好的归类
        "2025年2月20日至22日集装箱船的行驶轨迹？",
        "昨天有哪些集装箱船经过南京长江大桥卡口？",
        "本周有没有集装箱船存在违规行为？", # 这个问题无法很好的归类
        "长江主航道当前有哪些集装箱船？",
        "哪些区域禁止排放船舶生活污水",
        "船舶在狭水道如何通行",
        "船舶排放污染物要如何操作",
        "什么规格的船舶需要设置垃圾告示牌",
        "船舶垃圾有哪些分类",
        "违反《400 总吨以下内河船舶水污染防治管理办法》怎么处罚",
        "黄浦江的范围是什么",
        "船舶在杨浦大桥至徐浦大桥之间可以鸣笛吗",
        "黄浦江航道是单向航行吗",
        "在黄浦江航行的安全要求是什么",
        "黄浦江航行的船速最多是多少",
        ]

for example_question in test_queries:
  print('example_question: ',example_question)
  result = classify_and_extract(example_question)
  print(result)